import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, export_text, plot_tree
from sklearn.ensemble import RandomForestClassifier
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import silhouette_score, classification_report
from scipy.stats import ttest_ind, f_oneway
import sys
import os

# 设置中文显示
import matplotlib
matplotlib.rcParams['font.family'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False

# 创建输出目录
os.makedirs('B', exist_ok=True)

# 设置输出重定向
LOG_FILE = 'B/B.txt'
log_file = open(LOG_FILE, 'w', encoding='utf-8')
original_stdout = sys.stdout
sys.stdout = log_file

print("="*100)
print("古代玻璃制品成分分析与鉴别 - 问题2：高钾玻璃与铅钡玻璃分类及亚类划分")
print("="*100)
print("基于预处理.xlsx数据进行建模分析\n")

# 1. 数据加载与预处理
print("步骤1: 数据加载与预处理")
print("-"*100)
df = pd.read_excel('handled/预处理.xlsx', sheet_name='Sheet1')

print(f"原始数据维度: {df.shape}")
print(f"缺失值统计:\n{df.isnull().sum()}")

# 成分列列表
components = [
    '二氧化硅(SiO2)', '氧化钠(Na2O)', '氧化钾(K2O)', '氧化钙(CaO)', 
    '氧化镁(MgO)', '氧化铝(Al2O3)', '氧化铁(Fe2O3)', '氧化铜(CuO)', 
    '氧化铅(PbO)', '氧化钡(BaO)', '五氧化二磷(P2O5)', '氧化锶(SrO)', 
    '氧化锡(SnO2)', '二氧化硫(SO2)'
]

df[components] = df[components].fillna(0)
df['成分总和'] = df[components].sum(axis=1)
valid_data = df[(df['成分总和'] >= 85) & (df['成分总和'] <= 105)]
print(f"有效数据比例: {len(valid_data)/len(df):.2%} ({len(valid_data)}/{len(df)})")
df = valid_data.copy()

# 可视化1：数据质量概览
plt.figure(figsize=(15, 5))
plt.subplot(1, 3, 1)
plt.hist(df['成分总和'], bins=20, alpha=0.7, color='skyblue', edgecolor='black')
plt.title('成分总和分布')
plt.xlabel('成分总和(%)')
plt.ylabel('频数')

plt.subplot(1, 3, 2)
type_counts = df['类型'].value_counts()
plt.pie(type_counts.values, labels=type_counts.index, autopct='%1.1f%%')
plt.title('玻璃类型分布')

plt.subplot(1, 3, 3)
missing_data = df[components].isnull().sum()
top_missing = missing_data.head(8)
plt.bar(range(len(top_missing)), top_missing.values)
plt.xticks(range(len(top_missing)), top_missing.index, rotation=45)
plt.title('缺失值统计')
plt.ylabel('缺失数量')

plt.tight_layout()
plt.savefig('B/data_overview.png', bbox_inches='tight', dpi=300)
plt.close()
print("数据质量概览图已保存")

# 2. 主分类模型
print("\n步骤2: 高钾玻璃与铅钡玻璃主分类模型")
print("-"*100)

X = df[components]
y = df['类型']

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

print(f"训练集大小: {X_train.shape[0]}, 测试集大小: {X_test.shape[0]}")

# 训练模型
dt_classifier = DecisionTreeClassifier(
    max_depth=4, min_samples_split=5, min_samples_leaf=3, random_state=42
)
dt_classifier.fit(X_train, y_train)

rf_classifier = RandomForestClassifier(n_estimators=100, random_state=42)
rf_classifier.fit(X_train, y_train)

train_acc = dt_classifier.score(X_train, y_train)
test_acc = dt_classifier.score(X_test, y_test)
print(f"决策树训练集准确率: {train_acc:.4f}")
print(f"决策树测试集准确率: {test_acc:.4f}")

# 可视化2：模型性能与特征重要性
feature_importances = pd.DataFrame({
    '特征': components,
    '重要性': rf_classifier.feature_importances_
}).sort_values('重要性', ascending=False)

plt.figure(figsize=(15, 8))
plt.subplot(1, 2, 1)
plot_tree(dt_classifier, feature_names=components, class_names=['高钾', '铅钡'],
          filled=True, rounded=True, fontsize=8, max_depth=2)
plt.title('决策树可视化(前2层)')

plt.subplot(1, 2, 2)
top_features = feature_importances.head(10)
sns.barplot(x='重要性', y='特征', data=top_features, palette='viridis')
plt.title('特征重要性排序(前10)')
plt.xlabel('重要性')

plt.tight_layout()
plt.savefig('B/model_analysis.png', bbox_inches='tight', dpi=300)
plt.close()
print("模型分析图已保存")

# 3. 亚类划分
print("\n步骤3: 亚类划分")
print("-"*100)

# 3.1 铅钡玻璃亚类划分
lead_barium_df = df[df['类型'] == '铅钡'].copy()
print(f"铅钡玻璃样本数量: {len(lead_barium_df)}")

lead_barium_features = lead_barium_df[['氧化铅(PbO)', '氧化钡(BaO)', '二氧化硅(SiO2)', '氧化钾(K2O)']]
scaler = StandardScaler()
lead_barium_scaled = scaler.fit_transform(lead_barium_features)

# 轮廓系数评估
silhouette_scores = []
k_values = range(2, 6)
for k in k_values:
    kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
    cluster_labels = kmeans.fit_predict(lead_barium_scaled)
    silhouette_avg = silhouette_score(lead_barium_scaled, cluster_labels)
    silhouette_scores.append(silhouette_avg)
    print(f"K={k} : 轮廓系数 = {silhouette_avg:.4f}")

best_k = k_values[np.argmax(silhouette_scores)]
print(f"最佳聚类数: K={best_k}")

# 最终聚类
kmeans = KMeans(n_clusters=best_k, random_state=42, n_init=10)
lead_barium_df['亚类'] = kmeans.fit_predict(lead_barium_scaled) + 1

# 3.2 高钾玻璃亚类划分
potassium_df = df[df['类型'] == '高钾'].copy()
print(f"高钾玻璃样本数量: {len(potassium_df)}")

potassium_features = potassium_df[['氧化钾(K2O)', '二氧化硅(SiO2)', '氧化钙(CaO)', '氧化铝(Al2O3)']]
potassium_scaled = scaler.fit_transform(potassium_features)

# 确定高钾玻璃最佳聚类数
k_pot_scores = []
k_pot_values = range(2, 5)
for k in k_pot_values:
    kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
    cluster_labels = kmeans.fit_predict(potassium_scaled)
    silhouette_avg = silhouette_score(potassium_scaled, cluster_labels)
    k_pot_scores.append(silhouette_avg)

best_k_pot = k_pot_values[np.argmax(k_pot_scores)]
kmeans = KMeans(n_clusters=best_k_pot, random_state=42, n_init=10)
potassium_df['亚类'] = kmeans.fit_predict(potassium_scaled) + 1

# 可视化3：聚类分析
plt.figure(figsize=(15, 10))

# 轮廓系数比较
plt.subplot(2, 3, 1)
plt.plot(k_values, silhouette_scores, 'o-', label='铅钡玻璃')
plt.plot(k_pot_values, k_pot_scores, 's-', label='高钾玻璃')
plt.xlabel('聚类数K')
plt.ylabel('轮廓系数')
plt.title('最佳聚类数选择')
plt.legend()
plt.grid(True, alpha=0.3)

# 铅钡玻璃亚类分布
plt.subplot(2, 3, 2)
sns.scatterplot(x='氧化铅(PbO)', y='氧化钡(BaO)', hue='亚类', 
                data=lead_barium_df, palette='tab10', s=80)
plt.title('铅钡玻璃亚类分布')
plt.grid(True, alpha=0.3)

# 高钾玻璃亚类分布
plt.subplot(2, 3, 3)
sns.scatterplot(x='氧化钾(K2O)', y='二氧化硅(SiO2)', hue='亚类', 
                data=potassium_df, palette='Set2', s=80)
plt.title('高钾玻璃亚类分布')
plt.grid(True, alpha=0.3)

# 主成分对比箱线图
plt.subplot(2, 3, 4)
key_components_data = []
for comp in ['氧化钾(K2O)', '氧化铅(PbO)', '氧化钡(BaO)']:
    for glass_type in ['高钾', '铅钡']:
        values = df[df['类型'] == glass_type][comp]
        for val in values:
            key_components_data.append({'成分': comp, '类型': glass_type, '含量': val})

comp_df = pd.DataFrame(key_components_data)
sns.boxplot(data=comp_df, x='成分', y='含量', hue='类型')
plt.title('关键成分含量对比')
plt.xticks(rotation=45)

# 铅钡亚类成分均值热力图
plt.subplot(2, 3, 5)
lead_subgroup_means = lead_barium_df.groupby('亚类')[['氧化铅(PbO)', '氧化钡(BaO)', '二氧化硅(SiO2)', '氧化钾(K2O)']].mean()
sns.heatmap(lead_subgroup_means.T, annot=True, fmt='.2f', cmap='Blues')
plt.title('铅钡玻璃亚类成分均值')

# 高钾亚类成分均值热力图
plt.subplot(2, 3, 6)
pot_subgroup_means = potassium_df.groupby('亚类')[['氧化钾(K2O)', '二氧化硅(SiO2)', '氧化钙(CaO)', '氧化铝(Al2O3)']].mean()
sns.heatmap(pot_subgroup_means.T, annot=True, fmt='.2f', cmap='Greens')
plt.title('高钾玻璃亚类成分均值')

plt.tight_layout()
plt.savefig('B/clustering_analysis.png', bbox_inches='tight', dpi=300)
plt.close()
print("聚类分析图已保存")

# 4. 统计分析
print("\n步骤4: 统计分析")
print("-"*100)

# 主分类合理性分析
print("主分类合理性分析:")
key_stats = []
for comp in ['氧化钾(K2O)', '氧化铅(PbO)', '氧化钡(BaO)']:
    high_k = df[df['类型'] == '高钾'][comp]
    lead_ba = df[df['类型'] == '铅钡'][comp]
    
    t_stat, p_value = ttest_ind(high_k, lead_ba, equal_var=False)
    key_stats.append({
        '成分': comp,
        '高钾均值': high_k.mean(),
        '铅钡均值': lead_ba.mean(),
        'p值': p_value
    })
    print(f"{comp}: 高钾={high_k.mean():.2f}±{high_k.std():.2f}, 铅钡={lead_ba.mean():.2f}±{lead_ba.std():.2f}, p={p_value:.4f}")

# 亚类划分合理性分析
print("\n亚类划分合理性分析:")
print("铅钡玻璃亚类间差异:")
for comp in ['氧化铅(PbO)', '氧化钡(BaO)', '二氧化硅(SiO2)']:
    subgroups = [lead_barium_df[lead_barium_df['亚类'] == i][comp] for i in range(1, best_k+1)]
    f_stat, p_value = f_oneway(*subgroups)
    print(f"  {comp}: ANOVA p值 = {p_value:.6f}")

print("高钾玻璃亚类间差异:")
for comp in ['氧化钾(K2O)', '二氧化硅(SiO2)', '氧化钙(CaO)']:
    subgroups = [potassium_df[potassium_df['亚类'] == i][comp] for i in range(1, best_k_pot+1)]
    f_stat, p_value = f_oneway(*subgroups)
    print(f"  {comp}: ANOVA p值 = {p_value:.6f}")

# 可视化4：统计分析结果
plt.figure(figsize=(12, 8))

plt.subplot(2, 2, 1)
stats_df = pd.DataFrame(key_stats)
x_pos = np.arange(len(stats_df))
plt.bar(x_pos - 0.2, stats_df['高钾均值'], 0.4, label='高钾玻璃', alpha=0.8)
plt.bar(x_pos + 0.2, stats_df['铅钡均值'], 0.4, label='铅钡玻璃', alpha=0.8)
plt.xticks(x_pos, stats_df['成分'], rotation=45)
plt.ylabel('含量(%)')
plt.title('关键成分含量对比')
plt.legend()

plt.subplot(2, 2, 2)
perturbation_results = []
for feature in feature_importances['特征'].head(5):
    reduced_features = [f for f in components if f != feature]
    X_train_reduced = X_train[reduced_features]
    X_test_reduced = X_test[reduced_features]
    
    dt_reduced = DecisionTreeClassifier(max_depth=4, random_state=42)
    dt_reduced.fit(X_train_reduced, y_train)
    
    test_acc_reduced = dt_reduced.score(X_test_reduced, y_test)
    accuracy_drop = test_acc - test_acc_reduced
    perturbation_results.append(accuracy_drop)

plt.bar(range(5), perturbation_results)
plt.xticks(range(5), feature_importances['特征'].head(5), rotation=45)
plt.ylabel('准确率下降')
plt.title('特征重要性敏感性分析')

plt.subplot(2, 2, 3)
lead_cluster_sizes = lead_barium_df['亚类'].value_counts().sort_index()
plt.pie(lead_cluster_sizes.values, labels=[f'亚类{i}' for i in lead_cluster_sizes.index], 
        autopct='%1.1f%%')
plt.title('铅钡玻璃亚类样本分布')

plt.subplot(2, 2, 4)
pot_cluster_sizes = potassium_df['亚类'].value_counts().sort_index()
plt.pie(pot_cluster_sizes.values, labels=[f'亚类{i}' for i in pot_cluster_sizes.index], 
        autopct='%1.1f%%')
plt.title('高钾玻璃亚类样本分布')

plt.tight_layout()
plt.savefig('B/statistical_analysis.png', bbox_inches='tight', dpi=300)
plt.close()
print("统计分析图已保存")

# 5. 结果总结
print("\n步骤5: 结果总结")
print("-"*100)

print("铅钡玻璃亚类特征总结:")
for subgroup in range(1, best_k+1):
    subgroup_data = lead_barium_df[lead_barium_df['亚类'] == subgroup]
    print(f"亚类{subgroup} (样本数: {len(subgroup_data)})")
    print(f"  氧化铅(PbO): {subgroup_data['氧化铅(PbO)'].mean():.2f}±{subgroup_data['氧化铅(PbO)'].std():.2f}%")
    print(f"  氧化钡(BaO): {subgroup_data['氧化钡(BaO)'].mean():.2f}±{subgroup_data['氧化钡(BaO)'].std():.2f}%")

print("\n高钾玻璃亚类特征总结:")
for subgroup in range(1, best_k_pot+1):
    subgroup_data = potassium_df[potassium_df['亚类'] == subgroup]
    print(f"亚类{subgroup} (样本数: {len(subgroup_data)})")
    print(f"  氧化钾(K2O): {subgroup_data['氧化钾(K2O)'].mean():.2f}±{subgroup_data['氧化钾(K2O)'].std():.2f}%")
    print(f"  二氧化硅(SiO2): {subgroup_data['二氧化硅(SiO2)'].mean():.2f}±{subgroup_data['二氧化硅(SiO2)'].std():.2f}%")

# 关闭输出重定向
sys.stdout = original_stdout
log_file.close()

print(f"分析完成！详细结果已保存至 {LOG_FILE}")
print(f"可视化结果保存至 B/ 目录")
print("生成的图片文件:")
print("- B/data_overview.png: 数据质量概览")
print("- B/model_analysis.png: 模型分析与特征重要性")
print("- B/clustering_analysis.png: 聚类分析详情")
print("- B/statistical_analysis.png: 统计分析结果")